package com.wt.admin.service.vector.impl;


import com.wt.admin.service.vector.Vector;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.FileSystemResource;

import java.io.File;
import java.util.List;
import java.util.Map;

public class ESVectorImpl implements Vector {

    @Resource
    private VectorStore vectorStore;

    @Override
    public void add(String doc,Integer modelId,String FileId) {
        Document document = new Document(doc, Map.of("modelId",modelId,"fileId",FileId));
        TokenTextSplitter splitter = new TokenTextSplitter(30, 4, 10, 100,false);
        List<Document> apply = splitter.apply(List.of(document));
        vectorStore.add(apply);
    }

    @Override
    public void add(File path,Integer modelId,String FileId) {
        TikaDocumentReader tikaDocumentReader = new
                TikaDocumentReader(new FileSystemResource(path));
        List<Document> documents = tikaDocumentReader.get();
        documents.forEach(i -> {
            i.getMetadata().put("modelId",modelId);
            i.getMetadata().put("fileId",FileId);
        });
        TokenTextSplitter splitter = new TokenTextSplitter(30, 4, 10, 100,false);
        List<Document> apply = splitter.apply(documents);
        vectorStore.add(apply);
    }

    @Override
    public void add(List<Document> documents) {
        TokenTextSplitter splitter = new TokenTextSplitter(30, 4, 10, 100,false);
        List<Document> apply = splitter.apply(documents);
        vectorStore.add(apply);
    }

    @Override
    public void add(String doc, Map<String, Object> map) {

    }

    @Override
    public void add(File path, Map<String, Object> map) {

    }

    @Override
    public VectorStore getVectorStore() {
        return vectorStore;
    }

    @Override
    public void deleteByKnowledgeTitleId(Integer id) {
        SearchRequest request = SearchRequest.builder()
                .filterExpression("modelId == "+id)
                .topK(5)
                .build();
        List<Document> results = vectorStore.similaritySearch(request);
        vectorStore.delete(results.stream().map(Document::getId).toList());
    }

    @Override
    public void deleteByFileId(List<String> data) {
        data.forEach(i -> {
            SearchRequest request = SearchRequest.builder()
                    .filterExpression("fileId == '"+i+"'")
                    .topK(5)
                    .build();
            List<Document> results = vectorStore.similaritySearch(request);
            vectorStore.delete(results.stream().map(Document::getId).toList());
        });
    }

}
